# -*- coding: utf-8 -*-
"""
A股专用工具函数
包含沪深300成分股获取、北向资金接口、涨跌停状态检测等A股特有功能
"""

import pandas as pd
import numpy as np
from zipline.api import symbol, get_datetime

def get_csi300_constituents(context, data):
    """
    获取沪深300成分股（带流动性、波动率和基本面过滤）
    
    参数:
        context: 策略上下文
        data: 数据对象
        
    返回:
        list: 符合条件的股票列表
    """
    # 获取沪深300成分股
    csi300 = [s for s in context.index_components if s in context.security_list]
    
    # 流动性筛选 (20日平均成交额 > 10亿)
    high_liquidity = []
    for stock in csi300:
        avg_volume = data.history(stock, 'volume', 20, '1d').mean()
        avg_price = data.history(stock, 'close', 20, '1d').mean()
        turnover = avg_volume * avg_price
        if turnover > 1e9:  # 10亿人民币
            high_liquidity.append(stock)
    
    # 波动率过滤 (30日波动率 < 35%)
    low_volatility = []
    for stock in high_liquidity:
        returns = data.history(stock, 'close', 30, '1d').pct_change()
        volatility = returns.std() * np.sqrt(252)  # 年化波动率
        if volatility < 0.35:
            low_volatility.append(stock)
    
    # 基本面筛选 (PE < 25, PB < 2, ROE > 10%)
    fundamental_stocks = []
    for stock in low_volatility:
        pe = data.current(stock, 'pe_ratio')
        pb = data.current(stock, 'pb_ratio')
        roe = data.current(stock, 'roe')
        
        if pe < 25 and pb < 2 and roe > 0.1:
            fundamental_stocks.append(stock)
    
    return fundamental_stocks

def get_northbound_flow(stock, days=5):
    """
    获取北向资金数据（模拟接口）
    
    参数:
        stock: 股票代码
        days: 获取最近N天的数据
        
    返回:
        float: 最近N天北向资金净流入额
    """
    # 在实际使用中应替换为真实数据接口
    # 这里使用随机数据模拟
    np.random.seed(hash(stock) % 1000)
    return np.random.uniform(-5e7, 5e7)  # 随机生成-5000万到5000万的资金流

def is_price_limit(stock, data):
    """
    检测涨跌停状态
    
    参数:
        stock: 股票代码
        data: 数据对象
        
    返回:
        bool: True表示涨跌停（暂停交易），False表示正常交易
    """
    current_price = data.current(stock, 'close')
    prev_close = data.history(stock, 'close', 2, '1d')[0]
    
    # 计算涨跌停价格（A股±10%）
    upper_limit = prev_close * 1.1
    lower_limit = prev_close * 0.9
    
    # 检测是否触及涨跌停
    if current_price >= upper_limit or current_price <= lower_limit:
        return True
    return False

def apply_t1_restriction(context, stock):
    """
    应用T+1交易限制
    
    参数:
        context: 策略上下文
        stock: 股票代码
        
    返回:
        int: 可交易数量
    """
    position = context.portfolio.positions[stock]
    return position.closeable_amount